'''
breif:
    demo for Cogvideox-flash
    This demo is based on the Cogvideox-flash model and the ZhipuAI API,
    This demo provides user a class to generate and download video based on Cogvideox-flash model.
    For more infomation about Cogvideox-flash model and how to use it, please refer to https://bigmodel.cn/dev/api/videomodel/cogvideox
'''
from zhipuai import ZhipuAI
import json,requests,base64
content  =  open('API.json','r',encoding='utf-8').read()
api = json.loads(content)['api']
client = ZhipuAI(api_key=api)

class AI():
    def __init__(self):
        ...
    ###
    # @id: str (id of the video)
    # @return: str (status of the video) and str (url of the video)
    ###
    def GetStatus(self, id):
        try:
            response = client.videos.retrieve_videos_result(
            id=id
            )
            if response.task_status == 'SUCCESS':
                return response.task_status, response.video_result[0].url
            else:
                return response.task_status, None
        except Exception as e:
            raise e
    ###
    # @id: str (id of the video)
    # @return: str (download information)
    ###
    def download(self,id):
        try:
            result, url = self.GetStatus(id)
            if result =='SUCCESS':
                resp = requests.get(url)
                with open(f'{id}.mp4','wb') as f:
                    f.write(resp.content)
                return f"download success, file name is {id}.mp4"
            else:
                return f"download failed, status is {result}"
        except Exception as e:
            raise e
    ###
    # @prompt: str (discription of the video),Max length: 512 Tokens
    # @quality: str (chose from 'quality'、'speed', default: 'quality')
    # @with_audio: bool (whether to add audio to the video, default: False)
    # @size: str (chose from 720x480、1024x1024、1280x960、960x1280、1920x1080、1080x1920、2048x1080、3840x2160. default: 1920x1080), size must under 5MB
    # @duration: int (duration of the video, default: 10)
    # @fps: int (frames per second, default: 30)
    # @image: str (path of the image) (default: None) image must be one of the following formats: .png、.jpeg、.jpg
    # @return: str (url of the video) or download information
    ###
    def generate(self, prompt="比得兔开小汽车，游走在马路上，脸上的表情充满开心喜悦。", quality='quality',with_audio=False,size=None,duration=10, fps=30,image=None):
        if image != None:
            with open(image, "rb") as image_file:
                encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
        try:
            response = client.videos.generations(
            model = "cogvideox-flash",
            prompt= prompt,
            quality = quality,
            with_audio = with_audio,
            size = size,
            duration = duration,
            fps = fps,
            image_url = encoded_string if image != None else None
            )
            result = f"id:{response.id}\nvideo_result:{response.video_result}\nstatus:{response.task_status}\nrequest_id={response.request_id}"
            return result
        except Exception as e:
            raise Exception(f"Error generating image: {e}")
if __name__ == "__main__":
    ai = AI()
    ### demo of generate video with prompt and image     
    # result = ai.generate(prompt="风轻轻的吹拂着小猫",image='image.png')
    # print(result)
    # statu = ai.GetStatus('your-id')
    # print(statu)
    # result = ai.download('your-id')
    # print(result)
    ###